# pip install numpy xarray netCDF4 matplotlib plotly
# pip install "dask[complete]" -q
# pip install geopy ipywidgets folium -q
# pip install pycountry -q
# pip install shapely
# !jupyter nbextension enable --py widgetsnbextension
#!jupyter labextension install @jupyter-widgets/jupyterlab-manager

from ipywidgets import Layout, Dropdown, widgets
from IPython.display import display, clear_output, IFrame
import plotly.graph_objects as go
from functools import partial
import modules.n1_utilities as uti
import datetime
import numpy as np
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)

country_list = uti.read_json_to_sorted_dict('countries.json')
months = uti.read_json_to_dict('months.json')
timescales = uti.read_json_to_dict('timescales.json')
subset_area = None
index = None

selected = {
    "country": None,
    "adm1_subarea": None,
    "adm2_subarea": None,
    "timescale": None,
    "month": None,
    "year": None,
    "year_range": None
}

placeholders = {
    "country": "country...",
    "adm1_subarea": "adm1_subarea...",
    "adm2_subarea": "adm2_subarea...",
    "timescale": "timescale...",
    "month": "month...",
    "year": "year..."
}
uti.save_selection(placeholders)

# Custom style and layout for descriptions and dropdowns
style = {'description_width': '150px'}
dropdown_layout = Layout(width='400px', display='flex', justify_content='flex-end')
range_layout = Layout(width='400px')
btn_layout = Layout(width='400px')


# Dropdown for countries
country_names = [country['name'] for country in country_list]
country_selector = widgets.Dropdown(
    options=[placeholders['country']] + country_names,
    description='Select a country:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for subareas, initially empty
adm1_subarea_selector = widgets.Dropdown(
    options=[placeholders['adm1_subarea']],
    description='a subarea of first level:',
    style=style,
    layout=dropdown_layout
)

adm2_subarea_selector = widgets.Dropdown(
    options=[placeholders['adm2_subarea']],
    description='or of second level:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for timescales
timescale_selector = widgets.Dropdown(
    options=[placeholders['timescale']] + list(timescales.keys()),
    description='Select a timescale:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for months
month_selector = widgets.Dropdown(
    options=[placeholders['month']] + list(months.keys()),
    description='Select a month:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for years
current_year = datetime.datetime.now().year
years_options = [str(year) for year in range(1940, current_year + 1)]

year_selector = widgets.Dropdown(
    options=[placeholders['year']] + years_options,
    description='Select a year:',
    disabled=False,
    style=style,
    layout=dropdown_layout
)

# SelectionRangeSlider for years
year_range_selector = widgets.SelectionRangeSlider(
    options=years_options,
    index=(len(years_options) - 1, len(years_options) - 1),  # Start and end at the last
    description='Select the year range:',
    disabled=False,
    style=style,
    layout=range_layout
)

selectors = {
    "country" : country_selector,
    "adm1_subarea": adm1_subarea_selector,
    "adm2_subarea": adm2_subarea_selector,
    "timescale": timescale_selector,
    "month": month_selector,
    "year": year_selector,
    "year_range": year_range_selector    
}


month_widgets_btn = widgets.Button(
    description='Get data',
    disabled=False,
    button_style='info', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click me',
    icon='filter', # (FontAwesome names without the `fa-` prefix)month
    layout=btn_layout
)
month_widgets_btn.custom_name='month_widgets_btn'


year_widgets_btn = widgets.Button(
    description='Get data',
    disabled=False,
    button_style='info', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click me',
    icon='filter', # (FontAwesome names without the `fa-` prefix)
    layout=btn_layout
)
year_widgets_btn.custom_name='year_widgets_btn'

year_range_widgets_btn = widgets.Button(
    description='Get data',
    disabled=False,
    button_style='info', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click me',
    icon='filter', # (FontAwesome names without the `fa-` prefix)
    layout=btn_layout
)
year_range_widgets_btn.custom_name='year_range_widgets_btn'

# Output area for display updates
output_area = widgets.Output()

def setup_observers():
    """
    Sets up observers for UI widgets to handle interactions and updates dynamically in a graphical user interface.
    This function ensures that observers are only set once using a function attribute to track whether observers have
    already been established, enhancing efficiency and preventing multiple bindings to the same event.

    Observer is attached to widgets for country selection. This observer triggers specific functions when the 'value' property 
    of the widgets changes, facilitating responsive updates to the user interface
    based on user interactions.

    Notes:
    - This function uses a custom attribute `observers_set` on itself to ensure observers are set only once.
    """
    if not hasattr(setup_observers, 'observers_set'):      
            # When 'value' changes, update_subareas function will be called to update the dropdown menus
            # Create a partial function that includes the additional parameters
            country_selector.observe(partial(uti.update_subareas, 
                                         country_list=country_list, 
                                         placeholders=placeholders,
                                         adm1_subarea_selector=adm1_subarea_selector, 
                                         adm2_subarea_selector=adm2_subarea_selector), 'value')

            month_selector.observe(partial(uti.month_year_interaction, 
                                        month_selector=month_selector,
                                        year_selector=year_selector,
                                        selected=selected,
                                        placeholders=placeholders), 'value')
            year_selector.observe(partial(uti.month_year_interaction, 
                                        month_selector=month_selector,
                                        year_selector=year_selector,
                                        selected=selected,
                                        placeholders=placeholders), 'value')
            year_range_selector.observe(partial(uti.month_year_interaction, 
                                        month_selector=month_selector,
                                        year_selector=year_selector,
                                        selected=selected,
                                        placeholders=placeholders), 'value')
            # Set a flag to indicate observers are set
            setup_observers.observers_set = True



# Function to update and get data
def update_and_get_data(btn_name):
    global selected, placeholders, output_area, subset_data, index
    if uti.validate_selections(btn_name, selected, selectors, placeholders, output_area):
        with output_area:
            output_area.clear_output(wait=True)
            boundaries = uti.get_boundaries(selected, country_list, placeholders)
            coordinates = np.array(boundaries['features'][0]['geometry']['coordinates'][0])
            
            # Calculating the bounding box
            min_lon = np.min(coordinates[:, 0])
            max_lon = np.max(coordinates[:, 0])
            min_lat = np.min(coordinates[:, 1])
            max_lat = np.max(coordinates[:, 1])
            
            # rounding the coordinates to the nearest grid point
            adjusted_min_lon = uti.nearest_grid_point(min_lon)
            adjusted_max_lon = uti.nearest_grid_point(max_lon)
            adjusted_min_lat = uti.nearest_grid_point(min_lat)
            adjusted_max_lat = uti.nearest_grid_point(max_lat)

            bounding_box = (adjusted_min_lon, adjusted_min_lat, adjusted_max_lon, adjusted_max_lat)
            
            #print('Original Coordinates Sample: ', coordinates[:5])  # Showing first 5 for brevity
            #print('Bounding Box: ', bounding_box)
            
            # Fetching data using the bounding box
            subset_data = uti.get_xarray_data(btn_name, bounding_box, selectors, placeholders, months, timescales)
            index = f"SPEI{timescales[selectors['timescale'].value]}"
            adm_level, selected_area = uti.get_adm_level_and_area_name(selected, placeholders)
            timescale = selected['timescale']
            time_period = uti.get_period_of_time(btn_name, selected, placeholders)
                
            print(f"SPEI subset data uploaded for {selected_area}, administrative level {adm_level}, timescale {timescale}, period {time_period}")
            
            
# Set up widget interaction
def on_button_clicked(btn):
    update_and_get_data(btn.custom_name)


# Setup observers
setup_observers()

# Update existing selectors
previous_selection = uti.read_json_to_dict('selection.json')

# Set up widgets with previous settings
country_selector.value = previous_selection.get('country', placeholders['country'])
adm1_subarea_selector.value = previous_selection.get('adm1_subarea', placeholders['adm1_subarea'])
adm2_subarea_selector.value = previous_selection.get('adm2_subarea', placeholders['adm2_subarea'])
timescale_selector.value = previous_selection.get('timescale', placeholders['timescale'])
month_selector.value = previous_selection.get('month', placeholders['month'])
month_widgets_btn.on_click(on_button_clicked)

# Display widgets
display(country_selector, adm1_subarea_selector, adm2_subarea_selector, timescale_selector, month_selector, month_widgets_btn, output_area)

uti.display_data_details(selected, subset_data[index])

Country:  Afghanistan
ADM1 subarea:  adm1_subarea...
ADM2 subarea:  Acheen
Month:  February
Year:  year...
Year range:  ('2024', '2024')
Timescale:  6 months 

Time values in the subset: 85
Latitude values in the subset: 2
Longitude values in the subset: 2 

Data sample:  [[-9999. -9999.]
 [-9999. -9999.]]

processed_subset, change_summary = uti.process_datarray(subset_data[index])
print(processed_subset, '\n')
print('Change summary:')
for key, val in change_summary.items():
    print(key, val)

<xarray.DataArray 'SPEI6' (time: 85, lat: 2, lon: 2)>
dask.array<where, shape=(85, 2, 2), dtype=float64, chunksize=(1, 2, 2), chunktype=numpy.ndarray>
Coordinates:
  * time     (time) datetime64[ns] 1940-02-01T06:00:00 ... 2024-02-01T06:00:00
  * lon      (lon) float64 70.25 70.5
  * lat      (lat) float64 34.0 34.25
Attributes:
    long_name:  Standardized Drought Index (SPEI6)
    units:      - 

Change summary:
invalid_values_replaced 4
invalid_ratio 0.011764705882352941
duplicates_removed 0
cftime_conversions 1

stat_values = uti.compute_stats(processed_subset)

uti.create_scatterplot(stat_values, timescales, selected, placeholders)

uti.create_boxplot(stat_values, timescales, selected, placeholders)

# Update existing selectors
previous_selection = uti.read_json_to_dict('selection.json')

# Set up widgets with previous settings
country_selector.value = previous_selection.get('country', placeholders['country'])
adm1_subarea_selector.value = previous_selection.get('adm1_subarea', placeholders['adm1_subarea'])
adm2_subarea_selector.value = previous_selection.get('adm2_subarea', placeholders['adm2_subarea'])
timescale_selector.value = previous_selection.get('timescale', placeholders['timescale'])
year_selector.value = previous_selection.get('year', placeholders['year'])
year_widgets_btn.on_click(on_button_clicked)

# Display widgets
display(country_selector, adm1_subarea_selector, adm2_subarea_selector, timescale_selector, year_selector, year_widgets_btn, output_area)

uti.display_data_details(selected, subset_data[index])

Country:  Afghanistan
ADM1 subarea:  adm1_subarea...
ADM2 subarea:  Acheen
Month:  month...
Year:  1983
Year range:  ('2024', '2024')
Timescale:  6 months 

Time values in the subset: 12
Latitude values in the subset: 2
Longitude values in the subset: 2 

Data sample:  [[1.76448233 1.36480715]
 [1.3249789  0.76373523]]

processed_subset, change_summary = uti.process_datarray(subset_data[index])
print(processed_subset, '\n')
print('Change summary:')
for key, val in change_summary.items():
    print(key, val)

<xarray.DataArray 'SPEI6' (time: 12, lat: 2, lon: 2)>
dask.array<where, shape=(12, 2, 2), dtype=float64, chunksize=(1, 2, 2), chunktype=numpy.ndarray>
Coordinates:
  * time     (time) datetime64[ns] 1983-01-01T06:00:00 ... 1983-12-01T06:00:00
  * lon      (lon) float64 70.25 70.5
  * lat      (lat) float64 34.0 34.25
Attributes:
    long_name:  Standardized Drought Index (SPEI6)
    units:      - 

Change summary:
invalid_values_replaced 0
invalid_ratio 0.0
duplicates_removed 0
cftime_conversions 0

stat_values = uti.compute_stats(processed_subset)

uti.create_linechart(stat_values, timescales, selected, placeholders)

# Update existing selectors
previous_selection = uti.read_json_to_dict('selection.json')

# Set up widgets with previous settings
country_selector.value = previous_selection.get('country', placeholders['country'])
adm1_subarea_selector.value = previous_selection.get('adm1_subarea', placeholders['adm1_subarea'])
adm2_subarea_selector.value = previous_selection.get('adm2_subarea', placeholders['adm2_subarea'])
timescale_selector.value = previous_selection.get('timescale', placeholders['timescale'])
year_range_selector.value = previous_selection.get('year_range')
year_range_widgets_btn.on_click(on_button_clicked)

# Display widgets
display(country_selector, adm1_subarea_selector, adm2_subarea_selector, timescale_selector, year_range_selector, year_range_widgets_btn, output_area)

uti.display_data_details(selected, subset_data[index])

Country:  Afghanistan
ADM1 subarea:  adm1_subarea...
ADM2 subarea:  Acheen
Month:  month...
Year:  year...
Year range:  ('1940', '2024')
Timescale:  6 months 

Time values in the subset: 1012
Latitude values in the subset: 2
Longitude values in the subset: 2 

Data sample:  [[-9999. -9999.]
 [-9999. -9999.]]

processed_subset, change_summary = uti.process_datarray(subset_data[index])
print(processed_subset, '\n')
print('Change summary:')
for key, val in change_summary.items():
    print(key, val)

<xarray.DataArray 'SPEI6' (time: 1010, lat: 2, lon: 2)>
dask.array<getitem, shape=(1010, 2, 2), dtype=float64, chunksize=(1, 2, 2), chunktype=numpy.ndarray>
Coordinates:
  * time     (time) datetime64[ns] 1940-01-01T06:00:00 ... 2024-02-01T06:00:00
  * lon      (lon) float64 70.25 70.5
  * lat      (lat) float64 34.0 34.25
Attributes:
    long_name:  Standardized Drought Index (SPEI6)
    units:      - 

Change summary:
invalid_values_replaced 20
invalid_ratio 0.004940711462450593
duplicates_removed 2
cftime_conversions 4

stat_values = uti.compute_stats(processed_subset)

uti.create_stripechart(stat_values, timescales, selected, placeholders)

uti.create_stripechart(stat_values, timescales, selected, placeholders, 'year')

countries' boundaries source: www.geoboundaries.org